import pandas as pd
import time
import os
from weather.utils.utils import fix_seed_torch, draw_loss_pic
import argparse
from weather.utils.model import GCF
from weather.utils.mydataset import MyDataset
import torch


class Predictor:
    def __init__(self):
        fix_seed_torch(seed=2021)
        pre_args = {
            'gcn_layers': 2,
            'n_epochs': 100,
            'embedSize': 64,
            'batch_size': 1024,
            'lr': 0.001,
            'ratio': 0.8,
            'device': 'cpu'
        }

        self.base_dir = os.getcwd()

        self.user_feature = pd.read_csv(self.base_dir + '/weather/utils/data/user.txt', encoding='utf-8', sep='\t')
        self.item_feature = pd.read_csv(self.base_dir + '/weather/utils/data/weather.txt', encoding='utf-8', sep='\t')
        self.rating = pd.read_csv(self.base_dir + '/weather/utils/data/rating.txt', encoding='utf-8', sep='\t')
        self.model = GCF(pre_args, self.user_feature, self.item_feature, self.rating)

        self.model.load_state_dict(
            torch.load(self.base_dir + '/weather/utils/model/params.pth', map_location=torch.device('cpu')))

    def pred(self, user_id):
        user_id = 1
        item_num = self.rating['itemId'].max() + 1
        u = torch.tensor([int(user_id) for i in range(item_num)], dtype=torch.int)
        v = torch.tensor([i for i in range(item_num)], dtype=torch.int)
        result = self.model(u, v)  # 使用模型进行预测
        result = sorted(list(enumerate(result)), key=lambda x: x[1], reverse=True)[:5]
        print("用户最希望被推送的五种天气的ID是：".format(user_id))
        print([i[0] for i in result])
        return [i[0] for i in result]

    def pred_n(self, user_id: int, num: int):
        item_num = self.rating['itemId'].max() + 1
        u = torch.tensor([int(user_id) for i in range(item_num)], dtype=torch.int)
        v = torch.tensor([i for i in range(item_num)], dtype=torch.int)
        result = self.model(u, v)  # 使用模型进行预测
        result = sorted(list(enumerate(result)), key=lambda x: x[1], reverse=True)[:num]
        print("用户最希望被推送的{}种天气的ID是：{}".format(num, user_id))
        print([i[0] for i in result])
        return [i[0] for i in result]
